from sentence_transformers import CrossEncoder

def rerank(query: str, retrived_chunks: list[str], top_k: int) -> list[str]:
    cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
    pairs = [(query, chunk) for chunk in retrived_chunks]
    scores = cross_encoder.predict(pairs)

    chunk_with_score_list = [
        (chunk, score) for chunk, score in zip(retrived_chunks, scores)
    ]
    chunk_with_score_list.sort(key=lambda pair: pair[1], reverse=True)

    return [chunk for chunk,_ in chunk_with_score_list][:top_k]